import datetime as dt
import pandas as pd
import numpy as np
import cb_data as cb


def get_issueamount(code):
    return 1


def get_start_loc(obj, date):
    if date in obj.DB['Amt'].index:
        ret = obj.DB['Amt'].index.get_loc(date)
    else:
        # fake_index = pd.to_datetime(obj.DB['Amt'].index, format='%Y%m%d',errors ='raise')
        # fake_index = obj.DB['Amt'].index.values
        ret = obj.DB['Amt'].index.get_loc(date, method='bfill')

    return ret


def define_codes(obj, method=None):
    if not method or method == 'default':
        return obj.exclude_special()
    elif method == 'nonEB':
        return obj.exclude_special(hasEB=False)
    elif hasattr(method, '__call__'):
        return method(obj)


def select_codes(obj, codes, date, sel_method=None):
    i = get_start_loc(obj, date)
    n = min([i, 5])
    condition = (obj.DB['Amt'].iloc[i - n:i][codes].fillna(0).min() > 10000) & \
                (obj.DB['Outstanding'].iloc[i][codes] > 30000000)

    if sel_method:
        tempCodes = list(condition[condition].index)
        moreCon = sel_method(obj, codes, date, tempCodes)
        condition &= moreCon

        ret_codes = list(condition[condition].index)

        if not ret_codes:
            print('Empty selection on date: ' + date)

        return ret_codes


def get_weight(obj, codes, date, method=None):
    if not method or method == 'average':        # 等权策略
        # 这里要依赖一下numpy中的ones了
        ret = pd.Series(np.ones(len(codes)) / float(len(codes)), index=codes)
        return ret
    elif method == 'fakeEv':        # 按发行额加权，即“假市值”。中证转债指数类似这种
        srsIssue = get_issueamount(codes)
        srsFakeEv = obj.DB['Close'].loc[date, codes] * srsIssue
        return srsFakeEv / srsFakeEv.sum()
    elif method == 'Ev':    # 市值加权
        srsOutstanding = obj.DB['Outstanding'].loc[date, codes]
        srsEv = obj.DB['Close'].loc[date, codes] * srsOutstanding
        return srsEv / srsEv.sum()
    elif hasattr(method, '__call__'):
        return method(obj, codes, date)


def log_nav(obj, ret_df, asset_df, cash, date, cash_rate=0.03):
    if date == ret_df.index[0]:
        ret_df.loc[date]['NAV'] = 100
    else:
        i = ret_df.index.get_loc(date)
        j = obj.DB['Close'].index.get_loc(date)
        if len(asset_df.index) == 1 and asset_df.index[0] == 'Nothing':
            ret_df.iloc[i]['NAV'] = ret_df.iloc[i - 1]['NAV'] * (1 + cash_rate / 252.0)
            cash *= 1 + cash_rate / 252.0
        else:
            codes = list(asset_df.index)
            srsPct = obj.DB['Close'].iloc[j - 1:j + 1][codes].pct_change().iloc[-1] + 1.0
            cashW = 1 - asset_df['w'].sum()
            t1 = (srsPct * asset_df['costPrice'] * asset_df['w']).sum() + cash * cashW * (1 + cash_rate / 252)
            t0 = (asset_df['costPrice'] * asset_df['w']).sum() + cash * cashW
            ret_df.iloc[i]['NAV'] = ret_df.iloc[i - 1]['NAV'] * t1 / t0
            cash *= 1 + cash_rate / 252


def round_adjust(obj, start, method=None):
    i = get_start_loc(obj, start)
    if not method or method == 'daily':
        return obj.DB['Amt'].index[i:]
    elif isinstance(method, int):    # 这里有一个值得注意，验证数据类型，不要用 type(data) == ...，而是instance
        return obj.DB['Amt'].index[i:][::method]    # [::n]就是每隔n个数取一次了


def low_price(obj, codes, date, tempCodes):
    avg_price = obj.DB['Close'].loc[date][tempCodes].mean()
    if not avg_price:
        avg_price = 0
    return obj.DB['Close'].loc[date, codes] <= avg_price


def frame_strategy(obj, start='20170102',
                   define_method='default',
                   sel_method=low_price,
                   weight_method='average',
                   round_method=21):
    start_int = get_start_loc(obj, start)
    ret_df = pd.DataFrame(index=obj.DB['Amt'].index[start_int:], columns=['NAV', 'LOG:SEL', 'LOG:WEIGHT'])
    asset_df = pd.DataFrame(index=['Nothing'], columns=['costPrice', 'w'])
    cash = 100
    codes = define_codes(obj, define_method)
    adjust_date = round_adjust(obj, start, round_method)  # 调仓日期列表

    for i, date in enumerate(ret_df.index):
        log_nav(obj, ret_df, asset_df, cash, date)
        if date in adjust_date:
            sel = select_codes(obj, codes, date, sel_method)
            if sel:
                w = get_weight(obj, sel, date, weight_method)
            else:
                sel = ['Nothing']
                w = 0
            asset_df = pd.DataFrame(index=sel, columns=['costPrice', 'w'])
            asset_df['costPrice'] = 100
            asset_df['w'] = w

        ret_df['LOG:SEL'][date] = ','.join(list(asset_df.index))
        ret_df['LOG:WEIGHT'][date] = ','.join([str(t) for t in list(asset_df['w'])])

    return ret_df


if __name__ == '__main__':
    cb_obj = cb.CBData()
    cb_obj.load_data()
    result = frame_strategy(cb_obj)
